{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n\nCompile MXNet Models\n====================\n**Author**: `Joshua Z. Zhang <https://zhreshold.github.io/>`_,             `Kazutaka Morita <https://github.com/kazum>`_\n\nThis article is an introductory tutorial to deploy mxnet models with Relay.\n\nFor us to begin with, mxnet module is required to be installed.\n\nA quick solution is\n\n.. code-block:: bash\n\n    pip install mxnet --user\n\nor please refer to official installation guide.\nhttps://mxnet.apache.org/versions/master/install/index.html\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# some standard imports\nimport mxnet as mx\nimport tvm\nimport tvm.relay as relay\nimport numpy as np"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Download Resnet18 model from Gluon Model Zoo\n---------------------------------------------\nIn this section, we download a pretrained imagenet model and classify an image.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm.contrib.download import download_testdata\nfrom mxnet.gluon.model_zoo.vision import get_model\nfrom PIL import Image\nfrom matplotlib import pyplot as plt\n\nblock = get_model(\"resnet18_v1\", pretrained=True)\nimg_url = \"https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true\"\nimg_name = \"cat.png\"\nsynset_url = \"\".join(\n    [\n        \"https://gist.githubusercontent.com/zhreshold/\",\n        \"4d0b62f3d01426887599d4f7ede23ee5/raw/\",\n        \"596b27d23537e5a1b5751d2b0481ef172f58b539/\",\n        \"imagenet1000_clsid_to_human.txt\",\n    ]\n)\nsynset_name = \"imagenet1000_clsid_to_human.txt\"\nimg_path = download_testdata(img_url, \"cat.png\", module=\"data\")\nsynset_path = download_testdata(synset_url, synset_name, module=\"data\")\nwith open(synset_path) as f:\n    synset = eval(f.read())\nimage = Image.open(img_path).resize((224, 224))\nplt.imshow(image)\nplt.show()\n\n\ndef transform_image(image):\n    image = np.array(image) - np.array([123.0, 117.0, 104.0])\n    image /= np.array([58.395, 57.12, 57.375])\n    image = image.transpose((2, 0, 1))\n    image = image[np.newaxis, :]\n    return image\n\n\nx = transform_image(image)\nprint(\"x\", x.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compile the Graph\n-----------------\nNow we would like to port the Gluon model to a portable computational graph.\nIt's as easy as several lines.\nWe support MXNet static graph(symbol) and HybridBlock in mxnet.gluon\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "shape_dict = {\"data\": x.shape}\nmod, params = relay.frontend.from_mxnet(block, shape_dict)\n## we want a probability so add a softmax operator\nfunc = mod[\"main\"]\nfunc = relay.Function(func.params, relay.nn.softmax(func.body), None, func.type_params, func.attrs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "now compile the graph\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "target = \"cuda\"\nwith tvm.transform.PassContext(opt_level=3):\n    lib = relay.build(func, target, params=params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Execute the portable graph on TVM\n---------------------------------\nNow, we would like to reproduce the same forward computation using TVM.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm.contrib import graph_executor\n\ndev = tvm.cuda(0)\ndtype = \"float32\"\nm = graph_executor.GraphModule(lib[\"default\"](dev))\n# set inputs\nm.set_input(\"data\", tvm.nd.array(x.astype(dtype)))\n# execute\nm.run()\n# get outputs\ntvm_output = m.get_output(0)\ntop1 = np.argmax(tvm_output.numpy()[0])\nprint(\"TVM prediction top-1:\", top1, synset[top1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Use MXNet symbol with pretrained weights\n----------------------------------------\nMXNet often use `arg_params` and `aux_params` to store network parameters\nseparately, here we show how to use these weights with existing API\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def block2symbol(block):\n    data = mx.sym.Variable(\"data\")\n    sym = block(data)\n    args = {}\n    auxs = {}\n    for k, v in block.collect_params().items():\n        args[k] = mx.nd.array(v.data().asnumpy())\n    return sym, args, auxs\n\n\nmx_sym, args, auxs = block2symbol(block)\n# usually we would save/load it as checkpoint\nmx.model.save_checkpoint(\"resnet18_v1\", 0, mx_sym, args, auxs)\n# there are 'resnet18_v1-0000.params' and 'resnet18_v1-symbol.json' on disk"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "for a normal mxnet model, we start from here\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "mx_sym, args, auxs = mx.model.load_checkpoint(\"resnet18_v1\", 0)\n# now we use the same API to get Relay computation graph\nmod, relay_params = relay.frontend.from_mxnet(mx_sym, shape_dict, arg_params=args, aux_params=auxs)\n# repeat the same steps to run this model using TVM"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}